#!/usr/bin/env python3
# Author: Armit
# Create Time: 2022/11/20 

from argparse import ArgumentParser

from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import GaussianNB, CategoricalNB

from data import get_data, FEATURE_CAT, FEATURE_NUM
from utils import show_clf_metrics


def bayes(args):
  X, y = get_data(args.limit, FEATURE_NUM if args.type == 'num' else FEATURE_CAT)
  X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.7, random_state=42)
  print(f'dataset: {len(X_train)} for train. {len(X_test)} samples for test')

  if   args.type == 'num': model = GaussianNB()
  elif args.type == 'cat': model = CategoricalNB()
  else: raise ValueError

  model.fit(X_train, y_train)
  y_pred = model.predict(X_test)

  show_clf_metrics(y_test, y_pred)


if __name__ == '__main__':
  parser = ArgumentParser()
  parser.add_argument('-T', '--type', default='num', choices=['num', 'cat'], help='feature type')
  parser.add_argument('-N', '--limit', default=20000, type=int, help='limit dataset size')
  args = parser.parse_args()

  bayes(args)
